//
// Created by tianq on 24-11-27.
//
#pragma once
#include <algorithm>
#include <vector>
#include <random>

// see BFPRT73 paper
constexpr int groupSize = 5;

// restrict param type to numbers
template<typename T>
concept Number = std::is_arithmetic_v<T>;

/**
 * get nth element in given array.
 * brute forced, for comparison purpose
 * @param arr the array you want to determine its nth element
 * @param N [0,arr.size()-1], return nth smallest
 * @return nth smallest element
 */
template<Number num>
num BruteNth(std::vector<num> &arr, const unsigned long long N) {
    std::ranges::sort(arr);
    return arr[N];
}

/**
 * pick median of given array
 * @param arr the array to be parsed
 * @return median of given array, when arr size is odd, bigger one is chosen
 * @attention not a mathematical median when size of arr is even
 */
template<Number num>
num BruteMedian(std::vector<num> arr) {
    std::ranges::sort(arr);
    return arr[arr.size() / 2];
}

/**
 * pick first element from vector
 * @param arr the array to be parsed
 * @return first element from array
 */
template<Number num>
num PickFirst(const std::vector<num> &arr) {
    return arr.front();
}

/**
 * pick a random element from vector
 * @param arr the array to be parsed
 * @return random element from array
 */
template<Number num>
num PickRandom(const std::vector<num> &arr) {
    // #include "pcg_random.hpp"
    // pcg_extras::seed_seq_from<std::random_device> seed_source;
    // pcg32 rng(seed_source);
    std::random_device rd;
    std::mt19937_64 rng(rd());
    std::uniform_int_distribution dist(0, static_cast<int>(arr.size() - 1));
    return arr[dist(rng)];
}

/**
 * pick a likely median element from vector,used by BFPRT
 * @param arr the array to be parsed
 * @return element from array
 */
template<Number num>
num PickBfprt(const std::vector<num> &arr) {
    if (arr.size() <= groupSize) return BruteMedian(arr);

    auto it = arr.begin();
    std::vector<num> medians;
    while (it != arr.end()) {
        std::vector<num> group;
        while (group.size() < groupSize && it != arr.end()) {
            group.push_back(*it);
            ++it;
        }
        medians.push_back(BruteMedian(group));
    }

    return PickBfprt(medians);
}

/**
 * get nth element in given array
 * @param arr the array you want to determine its nth element
 * @param targetIndex [0,arr.size()-1], return nth smallest
 * @param PickPivot Pivot Picking strategy, PickRandom should be good enough
 * @return nth smallest element
 */
template<Number num>
num QuickSelect(std::vector<num> arr, const unsigned long long targetIndex,
                num (*PickPivot)(const std::vector<num> &) = PickBfprt) {
    // std::cout << std::endl << "targetIndex" << targetIndex << std::endl;
    // for (const num &num: arr) std::cout << num << ", ";
    // std::cout << std::endl;

    if (arr.size() == 1) return arr.front();

    const num pivot = PickPivot(arr);
    // std::cout << "pivot: " << pivot << std::endl;

    std::vector<num> left, mid, right;
    for (num &n: arr) {
        if (n < pivot) left.push_back(n);
        else if (n > pivot) right.push_back(n);
        else mid.push_back(n);
    }

    if (targetIndex < left.size()) {
        return QuickSelect(left, targetIndex, PickPivot);
    }
    if (targetIndex < left.size() + mid.size()) {
        return mid.front();
    }
    return QuickSelect(right, targetIndex - left.size() - mid.size(), PickPivot);
}

/*
 * usage demo
 * don't forget #include <iostream>

int main() {
    std::vector<int> arr = {2,3,4,8,5,6,5,3,4};
    for (const auto &num: arr) cout << num << ", ";
    std::cout << "BFPRT:\t" << QuickSelect(arr, 5, PickBfprt) << std::endl;
    std::cout << "First:\t" << QuickSelect(arr, 5, PickFirst) << std::endl;
    std::cout << "Random:\t" << QuickSelect(arr, 5, PickRandom) << std::endl;
    std::cout << "Brute:\t" << BruteNth(arr, 5) << std::endl;
    std::ranges::sort(arr);
    for (const auto &num: arr) std::cout << num << ", ";
}
*/
